import requests
import json
import re
from text_processor import TextVectorizer
from session_manager import session_manager

class RAGChain:
    def __init__(self, deepseek_api_key, vectorizer=None, vectorizer_params=None):
        """
        初始化RAG链
        
        参数:
        - deepseek_api_key: DeepSeek API密钥
        - vectorizer: TextVectorizer实例，如果为None则创建新实例
        - vectorizer_params: 创建TextVectorizer时的参数
        """
        self.deepseek_api_key = deepseek_api_key
        if vectorizer is not None:
            self.vectorizer = vectorizer
        else:
            # 默认参数
            params = vectorizer_params or {}
            self.vectorizer = TextVectorizer(**params)
        
        # 初始化查询扩展词表
        self.query_expansion = {
            "自然语言处理": ["NLP", "文本处理", "语言理解"],
            "机器学习": ["ML", "深度学习", "神经网络"],
            "人工智能": ["AI", "智能系统", "机器学习"]
            # 可以根据领域知识扩展更多同义词
        }

    def expand_query(self, query):
        """
        查询扩展，增加相关关键词以提高检索召回率
        
        参数:
        - query: 用户查询
        
        返回:
        - 扩展后的查询
        """
        # 检查是否为纯数字或包含电话号码的查询
        if re.search(r'\b\d{11}\b', query) or query.strip().isdigit():
            return query
        
        expanded_terms = [query]
        
        # 简单的同义词扩展
        for term, synonyms in self.query_expansion.items():
            if term in query:
                expanded_terms.extend(synonyms)
        
        # 移除重复项
        expanded_terms = list(set(expanded_terms))
        
        # 构建扩展查询
        if len(expanded_terms) > 1:
            expanded_query = query + " (也称为 " + ", ".join(expanded_terms[1:]) + ")"
            return expanded_query
        return query
    
    def retrieve_relevant_documents(self, query, n_results=3, similarity_metric="cosine"):
        """
        从TextVectorizer检索与查询相关的文档
        
        参数:
        - query: 用户查询
        - n_results: 要检索的结果数量
        - similarity_metric: 相似度度量方法("cosine"或"bm25")
        
        返回:
        - 包含文本和元数据的文档列表
        """
        # 对于电话号码查询，不扩展查询并强制使用余弦相似度
        if re.search(r'\b\d{11}\b', query):
            similarity_metric = "cosine"
            print(f"检测到电话号码查询，使用余弦相似度和原始查询")
            search_query = query
        else:
            # 非电话号码查询，进行扩展
            search_query = self.expand_query(query)
        
        # 使用TextVectorizer的query方法检索相关文档
        results = self.vectorizer.query(
            search_query,
            n_results=n_results,
            similarity_metric=similarity_metric
        )
        
        return results

    def build_prompt(self, query, documents, session_id=None, prompt_template=None):
        """
        构建提示词
        
        参数:
        - query: 用户查询
        - documents: 相关文档列表(包含文本和元数据)
        - session_id: 会话ID(可选)
        - prompt_template: 自定义提示模板(可选)
        
        返回:
        - 构建好的提示词
        """
        # 如果没有提供自定义模板，使用默认模板
        if prompt_template is None:
            # 优化提示模板，特别针对电话号码查询
            if re.search(r'\b\d{11}\b', query):
                prompt_template = "基于以下上下文信息，回答电话号码{query}属于哪家公司。如果上下文中明确包含该电话号码与公司的关联信息，请准确提取；如果没有相关信息，请明确说明无法回答。\n\n上下文:\n{context}\n\n回答:"
            else:
                prompt_template = "基于以下上下文信息，准确回答用户问题。如果上下文信息不足，请明确说明无法回答。\n\n上下文:\n{context}\n\n问题:{query}\n\n回答:"
        
        # 提取文档内容并按相关性排序
        document_texts = []
        for i, doc in enumerate(documents):
            if "text" in doc:
                # 添加文档相关性分数，帮助模型判断信息重要性
                text = f"文档 {i+1} (相关度: {doc.get('similarity', 0):.2f}): {doc['text']}"
                document_texts.append(text)
        
        context = "\n\n".join(document_texts)
        
        # 构建基础提示
        prompt = prompt_template.format(context=context, query=query)
        
        # 添加会话历史
        if session_id:
            history = session_manager.get_session_history(session_id)
            if history:
                prompt += "\n\n以下是之前的对话历史，有助于你更好地理解用户的问题:\n"
                for i, item in enumerate(history[-3:]):  # 只取最近3条历史
                    prompt += f"用户问题 {i+1}: {item['query']}\n"
                    prompt += f"助手回答 {i+1}: {item['answer']}\n\n"

        return prompt

    def call_deepseek_api(self, prompt, max_tokens=500):
        """
        调用DeepSeek API生成回答
        
        参数:
        - prompt: 提示词
        - max_tokens: 最大生成 tokens 数量
        
        返回:
        - 生成的回答
        """
        url = "https://api.deepseek.com/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.deepseek_api_key}"
        }
        data = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": "你是一个有用的助手。"},
                {"role": "user", "content": prompt}
            ],
            "max_tokens": max_tokens
        }
        
        try:
            response = requests.post(url, headers=headers, data=json.dumps(data))
            # 先检查是否为402错误
            if response.status_code == 402:
                return f"DeepSeek API调用失败: 余额不足或权限问题 (402 Payment Required)。请检查您的API密钥。"
            # 再检查其他HTTP错误
            elif response.status_code != 200:
                return f"DeepSeek API调用失败: HTTP状态码 {response.status_code}，{response.text}"
            # 如果状态码正常，解析响应
            result = response.json()
            return result["choices"][0]["message"]["content"]
        except Exception as e:
            return f"调用DeepSeek API失败: {str(e)}"

    def run(self, query, session_id=None, n_results=3, similarity_metric="bm25", use_reranking=False):
        """
        运行完整的RAG链
        
        参数:
        - query: 用户查询
        - session_id: 会话ID(可选)
        - n_results: 要检索的结果数量
        - similarity_metric: 相似度度量方法("cosine"或"bm25")
        - use_reranking: 是否使用重排序(基于LLM的相关性判断)
        
        返回:
        - 包含回答和引用文档的字典
        """
        # 1. 检索相关文档
        relevant_docs = self.retrieve_relevant_documents(
            query,
            n_results=n_results if not use_reranking else n_results * 2,  # 如果重排序，获取更多结果
            similarity_metric=similarity_metric
        )
        
        if not relevant_docs:
            return {
                "answer": "未找到相关文档，无法回答问题。",
                "references": []
            }
        
        # 2. 可选：使用LLM进行重排序
        if use_reranking and len(relevant_docs) > 1:
            relevant_docs = self._rerank_documents(query, relevant_docs, n_results)
        
        # 3. 构建提示词
        prompt = self.build_prompt(query, relevant_docs, session_id)
        
        # 4. 调用DeepSeek API
        answer = self.call_deepseek_api(prompt)
        
        # 5. 准备引用信息
        references = []
        for i, doc in enumerate(relevant_docs):
            if "text" in doc:
                # 提取文档元数据
                metadata = doc.get("metadata", {})
                source = metadata.get("source", f"文档 {i+1}")
                chunk_index = metadata.get("chunk_index", i)
                
                # 截取文档内容预览
                preview = doc["text"][:100] + ("..." if len(doc["text"]) > 100 else "")
                
                references.append({
                    "id": i+1,
                    "source": source,
                    "chunk_index": chunk_index,
                    "preview": preview,
                    "full_text": doc["text"],
                    "similarity": doc.get("similarity", 0)
                })
        
        # 6. 更新会话历史
        if session_id:
            # 检查会话是否存在，如果不存在则创建新会话
            if not session_manager.get_session(session_id):
                session_manager.create_session()
            # 更新会话历史，包含引用信息
            session_manager.update_session(session_id, query, answer, {
                "references": references
            })
        
        return {
            "answer": answer,
            "references": references
        }

    def _rerank_documents(self, query, documents, n_results):
        """
        使用LLM对检索到的文档进行重排序
        
        参数:
        - query: 用户查询
        - documents: 检索到的文档列表
        - n_results: 最终返回的文档数量
        
        返回:
        - 重排序后的文档列表
        """
        # 构建重排序提示
        prompt = "你需要根据与查询的相关性对以下文档进行重排序。\n\n查询: {query}\n\n文档列表:\n".format(query=query)
        
        for i, doc in enumerate(documents):
            prompt += f"文档 {i+1}: {doc['text'][:200]}...\n\n"
        
        prompt += "请按相关性从高到低返回文档编号列表，例如: [3, 1, 2]\n只返回数字列表，不要添加任何解释。"
        
        # 调用LLM进行重排序
        rerank_result = self.call_deepseek_api(prompt, max_tokens=100)
        
        try:
            # 解析LLM返回的列表
            import ast
            rerank_indices = ast.literal_eval(rerank_result)
            
            # 确保返回的是列表且元素是整数
            if not isinstance(rerank_indices, list) or not all(isinstance(idx, int) for idx in rerank_indices):
                raise ValueError("Invalid rerank format")
            
            # 调整索引（LLM返回的是1-based，我们需要0-based）
            rerank_indices = [idx - 1 for idx in rerank_indices if 0 < idx <= len(documents)]
            
            # 去重
            rerank_indices = list(dict.fromkeys(rerank_indices))
            
            # 取前n_results个结果
            rerank_indices = rerank_indices[:n_results]
            
            # 重新排序文档
            reranked_docs = [documents[idx] for idx in rerank_indices]
            return reranked_docs
        except Exception as e:
            print(f"Reranking failed: {e}")
            # 如果重排序失败，返回原始文档的前n_results个
            return documents[:n_results]

# 示例用法
if __name__ == "__main__":
    # 替换为你的DeepSeek API密钥
    api_key = "sk-a6d2da9cbeaf4cf59512a993de770cbd"
    rag_chain = RAGChain(deepseek_api_key=api_key)
    
    # 测试查询
    query = "什么是自然语言处理?"
    answer = rag_chain.run(query)
    print(f"问题: {query}")
    print(f"回答: {answer}")